"""
三维克里金插值：
    1.grid_c_data的输入是Grid_and_data.xlsx的路径
    2.coordinates_d的输入是Result_grid_coordinates.xlsx的路径
    3.自定义合适的输出路径
"""


import pandas as pd
import numpy as np
from skgstat import Variogram, OrdinaryKriging
from scipy.spatial import cKDTree
from concurrent.futures import ThreadPoolExecutor, as_completed
from tqdm import tqdm

# 读取网格 C 的数据
grid_c_data = pd.read_excel(
    "/home/ligongru/yuanzi/q2/Grid_and_data.xlsx",
    engine='openpyxl'
)

coordinates_c = grid_c_data[['x', 'y', 'z']].values.astype(float)  # 注意顺序为 X, Y, Z
stress_values_c = grid_c_data['S1'].values.astype(float)

# 创建 KD 树
tree = cKDTree(coordinates_c)

# 读取目标网格 D 的坐标
coordinates_d = pd.read_excel(
    "/home/ligongru/yuanzi/q2/Result_grid_coordinates.xlsx",
    engine='openpyxl'
)[['x', 'y', 'z']].values.astype(float)

# 定义处理单个点的函数
def process_point(input_point, tree, coordinates_c, stress_values_c):
    # 使用 KD 树查找与输入点距离在 100 以内的所有点
    indices = tree.query_ball_point(input_point, r=80)  # 查询半径为 100 的点的索引
    if not indices:  # 如果没有找到任何邻居点，则返回 NaN
        return np.nan
        # 限制邻居点数量z
    min_neighbors = 50  # 最小邻近点数
    max_neighbors = 100  # 最大邻近点数
    if len(indices) > max_neighbors:
        distances, indices = tree.query(input_point, k=max_neighbors) 

    # 根据找到的邻近点提取应力值
    neighbor_stresses = stress_values_c[indices]
    neighbor_coordinates = coordinates_c[indices]
     # 检查是否有足够的邻居点
    if len(neighbor_stresses) < min_neighbors:  # 确保至少有 30 个邻居点
        print(f"Warning: for point {input_point}, not enough neighbors were found within the range.")
        return np.nan  # 或者返回一个默认值
    
    # 创建变差函数模型
    V = Variogram(neighbor_coordinates, neighbor_stresses, model='spherical', normalize=False)

    # 构建 Kriging 模型
    ok = OrdinaryKriging(V)

    # 进行插值预测
    try:
        # 使用 transform 方法进行插值
        predicted_stress = ok.transform(input_point.reshape(1, -1))
    except Exception as e:
        # 记录错误并设置预测值为 NaN
        print(f"Error at point {input_point}: {e}")
        predicted_stress = np.nan

    return predicted_stress[0] if isinstance(predicted_stress, np.ndarray) else predicted_stress

# 使用 ThreadPoolExecutor 来并行处理每个点
max_workers = 20  # 你可以根据你的硬件条件调整这个数值
with ThreadPoolExecutor(max_workers=max_workers) as executor:
    # 提交所有任务
    futures = [executor.submit(process_point, point, tree, coordinates_c, stress_values_c) for point in coordinates_d]
    
    # 收集结果
    predicted_stresses = []
    for future in tqdm(as_completed(futures), total=len(futures), desc="Processing points"):
        predicted_stresses.append(future.result())

# 将结果保存到 DataFrame 并导出
results = pd.DataFrame({
    'X': coordinates_d[:, 0],
    'Y': coordinates_d[:, 1],
    'Z': coordinates_d[:, 2],
    'Predicted Stress S1': predicted_stresses
})

results.to_excel("/home/ligongru/yuanzi/q2/q2_result_spherical.xlsx", index=False)

print("所有预测完成，结果已保存。")